import argparse
import struct


def gen_size(file_size: int):
    edge_size = 256
    while True:
        if file_size + 4 <= edge_size * edge_size * 4:
            return edge_size
        else:
            edge_size += 16


def wrap_raw(raw_data: bytes):
    wrapped_data = struct.pack('<i', len(raw_data)) + raw_data
    edge_size = gen_size(len(wrapped_data))
    delta_size = edge_size * edge_size * 4 - len(wrapped_data)
    wrapped_data += b'\x00' * delta_size

    return wrapped_data, edge_size

def dump_raw(src: str):
    with open(src, 'rb') as f:
        data = f.read()
        wrapped_data, edge_size = wrap_raw(data)
        header = bytearray(b'BM6b\x8e\x00\x00\x00\x00\x006\x00\x00\x00(\x00\x00\x00p\x08\x00\x00\xa0\x05\x00\x00\x01\x00\x18\x00\x00\x00\x00\x00\x00b\x8e\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00')
        edge_bytes = struct.pack('<i', edge_size)
        # width - 18:22
        header[18], header[19], header[20], header[21] = edge_bytes[0], edge_bytes[1], edge_bytes[2], edge_bytes[3]
        # height - 22:26
        header[22], header[23], header[24], header[25] = edge_bytes[0], edge_bytes[1], edge_bytes[2], edge_bytes[3]
        header.extend(wrapped_data)

        return header
        

def load_raw(path: str):
    with open(path, 'rb') as f:
        data = f.read()
        size = struct.unpack('<i', data[54:58])[0]
        # print(f'load raw, size = {size}, {data[58:58+size]}')
        return data[58:58+size]

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', type=str)
    parser.add_argument('-o', type=str)
    parser.add_argument('-m', type=str)
    args = parser.parse_args()

    if args.m == 'dump':
        data = dump_raw(args.i)
        with open(args.o, 'wb') as f:
            f.write(data)
    elif args.m == 'load':
        data = load_raw(args.i)
        with open(args.o, 'wb') as f:
            f.write(data)


if __name__ == '__main__':
    main()